#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/6 11:37
@Author  : thezehui@gmail.com
@File    : 2.Runnable回退机制.py

TODO 由于LLMs的非确定性，使得具备优雅地处理错误的能力变得很重要。通过LCEL，可以轻松地为任何链添加Fallbacks
TODO 在LLM应用中会有许多可能的故障点，如LLM API问题、模型输出的质量差、亦或者其他集成问题等。那么，如何优雅的处理，并隔离这些问题呢？fallbacks将会提供这方面的能力。
TODO 更关键的是，fallbacks不仅可以应用在LLM级别，还可以应用在整个可运行级别。
TODO 在使用fallbacks时，为了避免出现无限重试的情况，建议将LLM包装器的重试功能关闭。这样，如果第一个包装器无法成功处理请求，它将立即失败，而不是无限重试。
"""
import dotenv
from langchain.chains import LLMChain
from langchain_community.chat_models.baidu_qianfan_endpoint import QianfanChatEndpoint
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

# 1.构建prompt与LLM，并将model切换为gpt-3.5-turbo-18k引发出错
prompt = ChatPromptTemplate.from_template("{query}")
# TODO 当使用一个模型不行时会切换到另一个模型执行
llm = ChatOpenAI(model="gpt-3.5-turbo-180k").with_fallbacks([ChatOpenAI()])
# llm = QianfanChatEndpoint().with_fallbacks([ChatOpenAI()])

# 2.构建链应用
chain = prompt | llm | StrOutputParser()

# 3.调用链并输出结果
content = chain.invoke({"query": "你好，你是?"})
print(content)
print("=======================================")

prompt1=ChatPromptTemplate.from_template("请为我写一首关于 {topic}的诗。")
# 定义一个主任务
main_task = prompt1 | ChatOpenAI(model="gpt-3.5-turbo-180k")

# 定义一个回退方案，例如使用备用模型或简化问题
prompt2=ChatPromptTemplate.from_template("请为我提供一个关于 {topic}的基本信息。")
fallback_task = prompt2 | ChatOpenAI()

# 将回退方案添加到主任务中 TODO 这里置顶一下报错回调
main_task = main_task.with_fallbacks([fallback_task])
response = main_task.invoke({"topic":"程序员"})
print(response.content)
